import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
import parl


# 灰度处理后的Model
class Model(parl.Model):
    def __init__(self, act_dim, obs_dim, dueling=True):
        super(Model, self).__init__()
        self.dueling = dueling
        self.obs_dim = obs_dim
        self.act_dim = act_dim

        feature_dim = 32

        self.fc1 = nn.Linear(32, feature_dim)
        self.fc2 = nn.Linear(feature_dim, feature_dim)
        self.fc3 = nn.Linear(feature_dim, act_dim)
        if dueling:
            self.fc4 = nn.Linear(feature_dim, 1)
        self.flatten = nn.Flatten()

    def forward(self, obs):
        out = self.flatten(obs)  # 将卷积后的输出展平
        if self.dueling:
            out = F.elu(self.fc1(out))
            out = F.elu(self.fc2(out))
            adv = self.fc3(out)
            value = self.fc4(out)
            Q = value + adv - paddle.mean(adv, axis=1, keepdim=True)
        else:
            out = F.elu(self.fc1(out))  # 全连接
            out = F.elu(self.fc2(out))  # 全连接
            Q = self.fc3(out)  # 全连接
        return Q


if __name__ == '__main__':
    model = Model(4, (200, 200))
    obs = np.zeros((32, 4, 200, 200), dtype=np.float32)
    obs = paddle.to_tensor(obs)
    Q = model(obs)
